package cn.kinoko.common.utils;

import lombok.extern.slf4j.Slf4j;

import java.time.Duration;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.LockSupport;

/**
 * @author kinoko
 */
@Slf4j
public class TokenBucketLimiter {
    // 桶的容量
    private final int capacity;
    // 当前令牌数量
    private final AtomicInteger tokens;
    // 每秒填充速率
    private final int rate;
    // 等待时间
    private final long waitTime;
    private final ScheduledExecutorService scheduler;

    public TokenBucketLimiter(int capacity, int rate, long waitTime) {
        this.capacity = capacity;
        this.tokens = new AtomicInteger(capacity);
        this.rate = rate;
        this.waitTime = waitTime;
        // 创建一个单线程的定时任务执行器，使用虚拟线程
        this.scheduler = Executors.newSingleThreadScheduledExecutor(Thread.ofVirtual().factory());
        // 初始化填充令牌的定时任务
        initRefillTask();
    }

    public TokenBucketLimiter(int capacity, int rate, long waitTime, ScheduledExecutorService scheduler) {
        this.capacity = capacity;
        this.tokens = new AtomicInteger(capacity);
        this.rate = rate;
        this.waitTime = waitTime;
        // 建议使用虚拟线程池
        this.scheduler = scheduler;
        // 初始化填充令牌的定时任务
        initRefillTask();
    }

    /**
     * 初始化填充令牌的定时任务
     */
    private void initRefillTask() {
        // 填充令牌的任务：如果令牌数量未达到桶的容量，成功增加令牌，跳出循环
        Runnable refillTask = () -> {
            // 如果令牌数量未达到桶的容量
            while (tokens.get() < capacity) {
                // CAS 操作，如果令牌数量未达到桶的容量，则成功增加令牌，跳出循环
                int currentTokens = tokens.get();
                if (tokens.compareAndSet(currentTokens, Math.min(currentTokens + rate, capacity))) {
                    // 成功增加令牌，跳出循环
                    break;
                }
            }
        };
        // 定时任务，每秒执行一次
        scheduler.scheduleAtFixedRate(refillTask, 0, 1, TimeUnit.SECONDS);
    }

    /**
     * 尝试获取令牌
     *
     * @return true 表示获取成功，false 表示获取失败
     */
    public boolean tryConsume() {
        boolean isRetry = false;
        while (true) {
            int currentTokens = tokens.get();
            // 无令牌可消费，等待令牌生成，重新尝试获取令牌，否则扣除令牌
            if (currentTokens == 0) {
                // 如果已经重试过，则返回false
                if (isRetry) return false;
                isRetry = true;
                // 睡指定时间
                LockSupport.parkNanos(Duration.ofMillis(waitTime).toNanos());
            } else {
                // 尝试减少令牌，成功则消费成功，否则继续尝试
                if (tokens.compareAndSet(currentTokens, Math.max(currentTokens - 1, 0))) {
                    return true;
                }
            }
        }
    }

    /**
     * 关闭定时任务，防止线程泄漏
     */
    public void shutdown() {
        scheduler.shutdown();
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucketLimiter tokenBucket = new TokenBucketLimiter(1, 1, 1000); // 桶容量为5，每秒填充1个令牌
        CountDownLatch latch = new CountDownLatch(10);
        for (int i = 0; i < 10; i++) {
            int finalI = i;
            CompletableFuture.runAsync(() -> {
                try {
                    if (tokenBucket.tryConsume()) {
                        log.info("Request {} is allowed.", finalI + 1);
                    } else {
                        log.info("Request {} is rejected.", finalI + 1);
                    }
                } finally {
                    latch.countDown();
                }
            });
        }
        latch.await();
        tokenBucket.shutdown();
    }
}
